#pragma once

#include "AccelerationStructure.h"
#include "AccelerationStructureGeometry.h"

class BottomLevelAccelerationStructure : public AccelerationStructure
{
public:

	struct BuilderDetails {
		VkBuildAccelerationStructureFlagsKHR flags;
		VkAccelerationStructureInstanceKHR accelerationStructureInstance;
		std::vector<RefCountPtr<AccelerationStructureGeometry>> geometries;
		BuilderDetails()
		{
			flags = VK_BUILD_ACCELERATION_STRUCTURE_PREFER_FAST_TRACE_BIT_KHR;

			accelerationStructureInstance.instanceCustomIndex = 0;
			accelerationStructureInstance.mask = 0xFF;
			accelerationStructureInstance.instanceShaderBindingTableRecordOffset = 0;
			accelerationStructureInstance.flags = VK_GEOMETRY_INSTANCE_TRIANGLE_FACING_CULL_DISABLE_BIT_KHR;

			accelerationStructureInstance.transform = {
				1.0f, 0.0f, 0.0f, 0.0f,
				0.0f, 1.0f, 0.0f, 0.0f,
				0.0f, 0.0f, 1.0f, 0.0f };
		}
	};

	class Builder : public BuilderBase<BuilderDetails> {
		friend struct BuilderDetails;
	public:
		Builder() noexcept;
		Builder(Builder const& rhs) noexcept;
		Builder(Builder&& rhs) noexcept;
		~Builder() noexcept;
		Builder& operator=(Builder const& rhs) noexcept;
		Builder& operator=(Builder&& rhs) noexcept;

		Builder& SetFlags(VkBuildAccelerationStructureFlagsKHR flags) noexcept;

		// Position of the instance
		Builder& SetTransform(glm::mat4 transform) noexcept;
		// gl_InstanceCustomIndexEXT
		Builder& SetInstanceCustomIndex(uint32 instanceCustomIndex) noexcept;
		//  Only be hit if rayMask & instance.mask != 0
		Builder& SetMask(uint32 mask) noexcept;
		// if 0 We will use the same hit group for all objects
		Builder& SetInstanceShaderBindingTableRecordOffset(uint32 instanceShaderBindingTableRecordOffset) noexcept;
		Builder& SetInstanceFlags(VkGeometryInstanceFlagsKHR flags) noexcept;

		Builder& AddGeometry(RefCountPtr<AccelerationStructureGeometry> geometry) noexcept;

		RefCountPtr<BottomLevelAccelerationStructure> Build();
		RefCountPtr<BottomLevelAccelerationStructure> Build(RenderProcess* process, const std::string& name);

	private:
		friend class BottomLevelAccelerationStructure;
	};

	BottomLevelAccelerationStructure(BuilderDetails& builderDetails);

	virtual void ReleaseResource() override;

	bool IsBuilded() const;

	VkAccelerationStructureInstanceKHR& GetAccelerationStructureInstance();

	std::vector<VkAccelerationStructureGeometryKHR>& GetVkGeometries();

	std::vector<uint32>& GetMaxPrimitiveCounts();

	std::vector<VkAccelerationStructureBuildRangeInfoKHR*>& GetRangeInfos();

	VkAccelerationStructureBuildSizesInfoKHR& GetSizeInfo();

	VkAccelerationStructureKHR& GetClearStructure();

	void SaveClearStructure();

	void DeleteClearStructure();

	void GetInstanceAddress();

	void AddGeometry(RefCountPtr<AccelerationStructureGeometry> newGeometry);

	void RemoveGeometry(RefCountPtr<AccelerationStructureGeometry> oldGeometry);

	void UpdateTransform(VkTransformMatrixKHR transform);

	void UpdateBlas(VkBuildAccelerationStructureFlagsKHR flags = VK_BUILD_ACCELERATION_STRUCTURE_FLAG_BITS_MAX_ENUM_KHR);

protected:

	std::vector<RefCountPtr<AccelerationStructureGeometry>> mGeometries;

	VkAccelerationStructureInstanceKHR mAccelerationStructureInstance;

	std::vector<VkAccelerationStructureGeometryKHR> mVkGeometries;

	std::vector<uint32> mMaxPrimitiveCounts;
	
	std::vector<VkAccelerationStructureBuildRangeInfoKHR*> mAccelerationStructureBuildRangeInfos;

	VkAccelerationStructureBuildSizesInfoKHR mAccelerationStructureBuildSizesInfo;

	VkBuffer mClearStructureBuffer;

	VkDeviceMemory mClearStructureMemory;

	VkAccelerationStructureKHR mClearStructure;

	bool mIsBuilded;

};